import joblib
import numpy as np

class MalwareDetector:
    def __init__(self, model_path):
        """加载预训练模型"""
        self.model = joblib.load(model_path)
        self.scaler = joblib.load(model_path.replace('model', 'scaler'))
    
    def preprocess_input(self, features):
        """输入预处理"""
        scaled_features = self.scaler.transform([features])
        return np.expand_dims(scaled_features, axis=2)  # 为CNN添加通道维度
    
    def predict(self, features):
        """执行预测"""
        processed = self.preprocess_input(features)
        prediction = self.model.predict(processed)
        return 'Malicious' if prediction[0] > 0.5 else 'Benign' 